[benchmark] add option to enable CompiledAutograd #1536
Conversation
Signed-off-by: Masaki Kozuki <[email protected]>
crcrpar
force-pushed
the
crpa/litgptbench_compiledautograd
branch
from
fa70057 to
b772a75
Compare
|
Have you investigated further why this option helps the torchtitan model code and why there are no improvements here? |
|
no, I haven't had enough bandwidth for it |
It looks like torchtitan uses compiled_autograd only for ddp and not fsdp. In this PR where they added the compiled_autograd option, they did it only for ddp. While, from the comments they seem confident that compiled autograd brings an advantage for ddp, they don't seem to be as optimistic for fsdp. Here implemented the ddp compiled_autograd option: https://github.com/pytorch/torchtitan/blob/49c6d6fc15ef644e5c3b1003ad4e0d9ea5fcb9a9/torchtitan/parallelisms/parallelize_llama.py#L105-L110 |
|
Does the all reduce operation appear in the FX graph with DDP + compiled_autograd? |
What does this PR do?
CompiledAutograd seems to speed up FSDP2 and I checked it with torchtitan.
I however somehow do not find it beneficial for litgpt models.
setting: pjnl-20241209, 8H100
torchtitan Llama-3-8B
This uses activation checkpointing, since the provided config by default uses it -- https://github.com/pytorch/torchtitan/blob/05a8b5e4c1de979c4b49ff36e6b09d6055db29b1/train_configs/llama3_8b.toml#L53-L55
litgpt llama-2-7b-hf